import meshio
import os
import traceback

gmsh_cell_type_map = {
    "vertex": 15,
    "triangle": 2,
    "quad": 3,
    "tetra": 4,
    "hexahedron": 5,
    "wedge": 6,
    "pyramid": 7,
    "tetra10": 11,
    "hexahedron20": 12,
    "wedge15": 13,
    "pyramid13": 14,
    #"polyhedron8":5,
    #"polyhedron10":11
    # 可根据需要扩展
}

def triangulate_face(face_vertices):
    """简单三角化，将多边形面拆成三角形"""
    tris = []
    for i in range(1, len(face_vertices) - 1):
        tris.append([face_vertices[0], face_vertices[i], face_vertices[i+1]])
    return tris

def write_basic_gmsh22(mesh: meshio.Mesh, filename: str):
    with open(filename, "w") as f:
        # Header
        f.write("$MeshFormat\n")
        f.write("2.2 0 8\n")
        f.write("$EndMeshFormat\n")

        # Points
        f.write("$Nodes\n")
        f.write(f"{len(mesh.points)}\n")
        for i, pt in enumerate(mesh.points):
            f.write(f"{i+1} {pt[0]} {pt[1]} {pt[2]}\n")
        f.write("$EndNodes\n")

        # Elements
        f.write("$Elements\n")
        element_id = 1
        all_elements = []

        for cell_block in mesh.cells:
            cell_type = cell_block.type
            if cell_type in ("polyhedron8", "polyhedron10"):  # 多面体单元类型
                # cell_block.data 是多面体单元点索引和面的组合，需要额外解析
                # meshio对于polyhedron类型，单元数据格式是 list of dict 或复杂格式，  
                # 这里假设cell_block.data是包含faces信息的结构
                # 实际需要看meshio版本或结构做对应调整
                # 假设 cell_block.data[i] 是一组面的顶点索引列表，例如：
                # [ [face1_vertices], [face2_vertices], ... ]
                # 以下演示遍历每个单元的每个面，三角化，写三角形单元

                for polyhedron_faces in cell_block.data:
                    # polyhedron_faces: 多面体的所有face顶点索引列表（list of list）
                    for face in polyhedron_faces:
                        triangles = triangulate_face(face)
                        for tri in triangles:
                            # tri是三角形的3个点索引
                            node_str = " ".join(str(n + 1) for n in tri)  # Gmsh节点从1开始
                            line = f"{element_id} 2 2 0 0 {node_str}\n"  # 2是三角形单元类型
                            all_elements.append(line)
                            element_id += 1
            else:
                if cell_type not in gmsh_cell_type_map:
                    print(f"[WARN] Skip unsupported cell type: {cell_type}")
                    continue
                gmsh_type = gmsh_cell_type_map[cell_type]
                for cell in cell_block.data:
                    node_str = " ".join(str(pt_id + 1) for pt_id in cell)
                    line = f"{element_id} {gmsh_type} 2 0 0 {node_str}\n"
                    all_elements.append(line)
                    element_id += 1

        f.write(f"{len(all_elements)}\n")
        try:
            for line in all_elements:
                f.write(line)
            f.write("$EndElements\n")
        except Exception as e:
            traceback.print_exc()
            
# vtk_path：vtk文件路径
# gmsh_path：输出的gmsh文件路径
def VtkToGmsh(vtk_path: str, gmsh_path: str):
    # 判断文件是否存在
    is_exist = os.path.exists(vtk_path)
    if not is_exist:
        print("[ERROR] Vtk(" + vtk_path + ") is not exist.")

    try:
        mesh = meshio.read(vtk_path)
        write_basic_gmsh22(mesh,gmsh_path)
    except Exception as e:
        print("[ERROR] Vtk(" + vtk_path + ") to gmsh(" + gmsh_path + ") failed.")
        traceback.print_exc()


# gmsh_dir：gmsh输出文件夹路径
def MultiVtkToGmsh(vtk_paths: list[str], gmsh_dir: str):
    for vtk_path in vtk_paths:
        gmsh_path = ""
        try:
            # vtk文件名
            vtk_file_name = os.path.basename(vtk_path)
            file_name, ext = os.path.splitext(vtk_file_name)

            gmsh_file_name = file_name + ".msh"
            gmsh_path = gmsh_dir + "/" + gmsh_file_name

            VtkToGmsh(vtk_path, gmsh_path)
        except:
            print("[ERROR] Vtk(" + vtk_path + ") to gmsh(" + gmsh_path + ") failed.")
